# encoding:utf-8
'''
Created on 2020-12-23
@author: adog
ref: https://scikit-learn.org/stable/auto_examples/applications/plot_stock_market.html#sphx-glr-auto-examples-applications-plot-stock-market-py
'''
# import apps.adog.dataMan as DM
import util.xtool as xtool
import util.tradeDate as TD
# import apps.adog.util.dbcom as dbcom
import util.mdbcom as mdbcom

import numpy as np
import pandas as pd
# import json
import chart_studio.plotly
import plotly.graph_objs as graphObj

# import plotly.graph_objects as go
# from datetime import datetime, timedelta
# import time
import re
import math
import os 

import apps.adog.util.mibian as MIBIAN


class Calu():

    def __init__(self, cod=None, optCod=None, broker=None, cp=None):
        self.cod = cod
        self.broker = broker
        self.cp = cp
        self.opt_contr_info = None
        self.hv = None
        self.lastOptData = None
        
    def loadOptContInfo(self):
        self.opt_contr_info = opt.run_query(query(opt.OPT_CONTRACT_INFO).filter((opt.OPT_CONTRACT_INFO.code.like(self.cod + '%%'))\
                                                                               & (opt.OPT_CONTRACT_INFO.expire_date > xtool.nowDate())\
                                                                               & (opt.OPT_CONTRACT_INFO.contract_type == (self.cp + 'O'))).limit(300))
        self.opt_contr_info['dayToExp'] = self.opt_contr_info.apply(lambda df: xtool.deltaDays(xtool.nowDate(), str(df['expire_date'])), axis=1)        
        self.opt_contr_info['maturity'] = self.opt_contr_info['dayToExp'] / 365.
    
    def loadHV(self):
        self.hv = self.getHistVolatilityYangZhang(cod=self.cod, beginDate='2020-01-01', endDate=xtool.nowDate())
    
    def loadIV(self, optCodList, ivNormalised=None):
        # obsTime=xtool.dateBefAftX(1)
        obsTime = xtool.nowDate()
        if xtool.nowOnlyTimeMins() > '20:55':
            obsTime = TD.tradeDateBefAftX(obsTime, +1)
        all = get_price(optCodList, end_date=obsTime, count=1, frequency='1d', fields=['close'], panel=False)  # by default, time close included already!
        
        if len(all) > 0:
            
            all = pd.merge(left=all, right=self.opt_contr_info[['code', 'underlying_symbol', 'exercise_price', 'dayToExp', 'maturity']], on='code', how='left')  # by default, time close included already!
            underlying_symbol = all['underlying_symbol'].unique().tolist()
            fut = get_price(underlying_symbol, end_date=obsTime, count=1, frequency='1d', fields=['close'], panel=False)  # by default, time close included already!
            fut.rename(columns={
                                'time':'fut_time',
                                'code':'underlying_symbol',
                                'close':'underlyingPrice'
                                }, inplace=True)
            all = pd.merge(left=all, right=fut, on='underlying_symbol', how='left')
            all['iv'] = all.apply(lambda df:self.getOptRisk(underlyingPrice=df['underlyingPrice'],
                                                            strikePrice=df['exercise_price'],
                                                            interestRate=4.,
                                                            dayToExp=df['dayToExp'],
                                                            volatility=None,
                                                            vol_type='hv1W',
                                                            cp=self.cp,
                                                            opt_price=df['close'],
                                                            isOnlyCaluIV=True), axis=1)
            
            all['contractType'] = self.cp
            div = self.drawIVCurve(all, ivNormalised=ivNormalised)
            return div
        
    def drawIVCurve(self, df, ivNormalised=None):
        traceList_c = []
        traceList_p = []
        traceList_all = []
        # df['contrMonth']=df.apply(lambda d: re.sub('\D','', d['underlying_symbol']), axis=1)
        df['contrMonth'] = [re.sub('\D', '', df.loc[i, 'underlying_symbol']) for i in df.index]
        df_pivot = df.pivot(index='exercise_price', columns='contrMonth')
        
        ''' old version
        iv_atm=iv[np.argmin(np.abs(strikePrice-spotPrice))]
        yy=iv-iv_atm
        xx=np.log(strikePrice/spotPrice)/np.sqrt(t)
        '''
        df_iv = df_pivot['iv'].interpolate(axis=0).ffill(axis=0).bfill(axis=0)
        xx, yy = None, None
        if ivNormalised:
            
            iv_atm_index = np.argmin(np.abs(df_pivot['underlyingPrice'].ffill(axis=0).bfill(axis=0).values - df_pivot.index[:, None]).values, axis=0)
            # iv_atm_index2=[[iv_atm_index[i],i] for i in range(len(iv_atm_index))]
            iv_atm = [df_iv.values[iv_atm_index[i], i] for i in range(len(iv_atm_index))]
            yy = df_iv - iv_atm
            xx = np.log(df_pivot.index[:, None].values / df_pivot['underlyingPrice'].ffill(axis=0).bfill(axis=0).values) / np.sqrt(df_pivot['maturity'].ffill(axis=0).bfill(axis=0).values)
        else:
            yy = df_iv
            xx = df_pivot.index[:, None].values.repeat(len(df_pivot['iv'].columns), axis=1)
            
        for ci in range(len(yy.columns)):
            trace = graphObj.Scatter(
                        x=xx[:, ci],
                        y=yy.loc[:, yy.columns[ci]].values,
                        mode='lines+markers',
                        name=yy.columns[ci],
                        hoverinfo='name,y',
                        line=dict(
                            shape='dot'
                        )
                        )
            
            traceList_all.append(trace)  
            
        layout = dict(
            title='%s,ivNormalised=%s' % (self.cp, ivNormalised),
            height=600,
            width=800,
            legend=dict(
                y=0.5,
                traceorder='reversed',
                font=dict(
                    size=16
                )
            )
        ) 
        
        fig = dict(data=traceList_all, layout=layout)
        # plotly.offline.plot(fig, auto_open=True,filename='test.html')   
        # div_norm = plotly.offline.plot(fig, output_type='div',include_plotlyjs=False)  
        div = fig.to_html(full_html=True, include_plotlyjs=True, div_id='div_1')
        return div
        
    def loadLastOptData(self, optCod):
        # data=get_ticks(optCod, start_dt=xtool.nowTime(),end_dt=xtool.nowTime())
        # now=xtool.nowDate()+' '+xtool.nowOnlyTimeMins()+':00'
        # obsTime=xtool.dateBefAftX(1)
        
        obsDate = xtool.nowDate()
        if xtool.nowOnlyTimeMins() > '20:55':
            obsDate = TD.tradeDateBefAftX(obsDate, +1)
        
        lastPrice_opt = get_price(optCod, end_date=obsDate, count=5, frequency='1d')
        if len(lastPrice_opt) > 0:
            lastPrice_opt = lastPrice_opt.tail(1)['close'].values[0]
        info = self.opt_contr_info.loc[self.opt_contr_info.code == optCod, ['underlying_symbol', 'exercise_price', 'dayToExp']]
        underlying_symbol = 0.
        lastPrice_fur = 0.
        exercise_price = 0.
        dayToExp = 0.
        if len(info) == 1:
            underlying_symbol = info['underlying_symbol'].values[0]
            lastPrice_fur = get_price(underlying_symbol, end_date=obsDate, count=5, frequency='1d')
            if len(lastPrice_fur) > 0:
                lastPrice_fur = lastPrice_fur.tail(1)['close'].values[0]
            exercise_price = info['exercise_price'].values[0]
            dayToExp = info['dayToExp'].values[0]
        self.lastOptData = {
                'underlying_symbol':underlying_symbol,
                'underlying_price':lastPrice_fur,
                'opt_symbol':optCod,
                'lastPrice_opt':lastPrice_opt,
                'strikePrice':exercise_price,
                'dayToExp':dayToExp,
                'interestRate':4
                }
        
    def getOptRisk(self,
                       underlyingPrice=None,
                       strikePrice=None,
                       interestRate=None,  #  3
                       dayToExp=None,  # 12, day
                       volatility=None,
                       vol_type='hv1W',
                       cp=None,
                       opt_price=None,
                       isOnlyCaluIV=False
                       ):
        theoPrice = -99.
        iv = -99.
        '''
        BS([underlyingPrice, strikePrice, interestRate, daysToExpiration], \
        volatility=x, callPrice=y, putPrice=z)
        '''
        if not isOnlyCaluIV:
            if volatility is None:
                hv = self.getHistVolatilityYangZhang(cod=self.cod, beginDate='2020-01-01', endDate=xtool.nowDate())
                vola = self.hv.head(1)[vol_type]
                if len(vola) == 1:
                    volatility = vola.values[0]
                else:
                    volatility = 0
        
        if cp == 'C':
            if not isOnlyCaluIV:
                theoPrice = MIBIAN.BS([underlyingPrice, strikePrice, interestRate, dayToExp], volatility=volatility).callPrice
            iv = MIBIAN.BS([underlyingPrice, strikePrice, interestRate, dayToExp], callPrice=opt_price).impliedVolatility
        elif cp == 'P':
            if not isOnlyCaluIV:
                theoPrice = MIBIAN.BS([underlyingPrice, strikePrice, interestRate, dayToExp], volatility=volatility).putPrice
            iv = MIBIAN.BS([underlyingPrice, strikePrice, interestRate, dayToExp], putPrice=opt_price).impliedVolatility
        if theoPrice is None:
            theoPrice = -99.0    
        if iv is None:
            iv = -99.0 
            
        if isOnlyCaluIV:
            return round(iv, 2)
        else:
            return round(theoPrice, 2), round(iv, 2)
    
    def getHistVolatilityYangZhang(self, cod=None, beginDate=None, endDate=None):
        hv = None
        if xtool.nowOnlyTimeMins() > '20:55':
            endDate = TD.tradeDateBefAftX(endDate, +1)
        data = mdbcom.find('jq_hq_1d',
                         queryDict={'cod':cod.upper(),
                                   'code_type':'fin8',
                                   '$and':[{'date':{'$gte':beginDate}},
                                           {'date':{'$lte':endDate}}]   
                                   },
                         sortListOfTuple=[('date', 1)],
                         outType='pandas')
        if data is not None:
            data.set_index(np.arange(data.shape[0]), inplace=True) 
            pcNP = data['close'].values[:-1]
            pcList = [data.loc[0, 'close']] + pcNP.tolist()
            data['preClose'] = pcList
            # O
            data['sigmaO'] = np.log(data['open'] / data['preClose'])
            # C
            data['sigmaC'] = np.log(data['close'] / data['open'])
            # RS
            data['amp1'] = data['high'] / data['close']
            data['amp2'] = data['high'] / data['open']    
            data['amp3'] = data['low'] / data['close']
            data['amp4'] = data['low'] / data['open']    
            data['sigmaRS'] = np.log(data['amp1']) * np.log(data['amp2']) + np.log(data['amp3']) * np.log(data['amp4'])
            
            data['tradeDate'] = pd.to_datetime(data['date'])
            periods = {'hv1W': 5, 'hv2W': 10, 'hv1M': 21, 'hv2M': 41, 'hv3M': 62, 'hv4M': 83,
                       'hv5M': 104, 'hv6M': 124, 'hv9M': 186, 'hv1Y': 249, 'hv2Y': 497}
            # 利用Parkinson模型计算波动率
            for prd in periods.keys():
                sigmaO2 = np.power(data['sigmaO'].rolling(window=periods[prd]).std(), 2)
                sigmaC2 = np.power(data['sigmaC'].rolling(window=periods[prd]).std(), 2)
                sigmaRS2 = data['sigmaRS'].rolling(window=periods[prd]).mean()
                k = 0.34 / (1.34 + (periods[prd] + 1) / (periods[prd] - 1))
                sigma2 = np.sqrt(sigmaO2 + k * sigmaC2 + (1 - k) * sigmaRS2) * math.sqrt(252.0)
                data[prd] = np.round(sigma2 * 100, 2)
            
            data = data[data.date >= beginDate]
            data['xtime'] = xtool.nowTime()
            data['method'] = 'yz'
            fields = [u'_id',
                    u'cod',
                    u'code',
                    u'code_cn',
                    u'broker',
                    u'date',
                    u'tradeDate',
                    u'open',
                    u'high',
                    u'low',
                    u'close',
                    u'hv1W',
                    u'hv2W',
                    u'hv1M',
                    u'hv2M',
                    u'hv3M',
                    # u'hv4M',      
                    # u'hv5M',
                    u'hv6M',
                    u'hv1Y',
                    # u'hv2Y',      
                    u'xtime']
            hv = data[fields].copy()
            hv = hv[hv['date'].isin(TD.tradeDate)]
            hv.sort_values(by='date', ascending=0, inplace=True)
            hv.reset_index(drop=True, inplace=True)
        return hv
    
    def drawHv(self, hv):
        traceList_all = []
        fig = graphObj.Figure()
        for codcn in hv['cod_cn'].unique():
            hvx = hv.loc[hv['cod_cn'] == codcn,:]
            
            for ci in hvx.columns:
                if 'hv' in ci:
                    trace = graphObj.Scatter(
                                x=hvx.loc[:, 'date'].values,
                                y=hvx.loc[:, ci].values,
                                mode='lines',
                                name=ci + '_' + codcn,
                                hoverinfo='name,y',
                               
                                )
            
                    # fig.add_trace(trace)
            
                    traceList_all.append(trace)  
            if 'open' in hvx.columns:
                trace_fin8 = graphObj.Candlestick(x=hvx.loc[:, 'date'].values,
                                                open=hvx.loc[:, 'open'].values,
                                                high=hvx.loc[:, 'high'].values,
                                                low=hvx.loc[:, 'low'].values,
                                                close=hvx.loc[:, 'close'].values,
                                                yaxis='y2' ,
                                                name=codcn + '指数',
                                                increasing=dict(fillcolor='rgba(207,91,68,0.5)',
                                                                line=dict(color='rgb(212,98,75)', width=2)
                                                                ),
                                                decreasing=dict(fillcolor='rgba(77,220,85,0.5)',
                                                                line=dict(color='rgb(72,169,119)', width=2)
                                                                ),
                                                )
                # fig.add_trace(trace_fin8)
                traceList_all.append(trace_fin8)
        layout = dict(
            title='',
            height=900,
            width=2000,
            legend=dict(
                y=0.5,
                traceorder='reversed',
                font=dict(
                    size=16
                )
            ),
            yaxis2=dict(
                title='fin8',
                titlefont=dict(
                    color='rgb(148, 103, 189)'
                ),
                tickfont=dict(
                    color='rgb(148, 103, 189)'
                ),
                overlaying='y',
                side='right',
                # range=[-hv['high'].values.max()/2,hv['high'].values.max()]
                ) if 'open' in hvx.columns else dict()
        ) 
        
        fig.add_traces(traceList_all)
        fig.update_layout(title='',
                        height=900,
                        width=2000,
                        legend=dict(
                            y=0.5,
                            traceorder='reversed',
                            font=dict(
                                size=16
                            )
                        ),
                        yaxis2=dict(
                            title='fin8',
                            titlefont=dict(
                                color='rgb(148, 103, 189)'
                            ),
                            tickfont=dict(
                                color='rgb(148, 103, 189)'
                            ),
                            overlaying='y',
                            side='right',
                            # range=[-hv['high'].values.max()/2,hv['high'].values.max()]
                            ) if 'open' in hvx.columns else dict())
        # plotly.offline.plot(fig, auto_open=True,filename='test.html')   
        div = fig.to_html(full_html=True, include_plotlyjs=True, div_id='div_3')
        
        return div

        
if __name__ == '__main__':
    x = Calu(cod='c')
    p = x.getOptRisk(underlyingPrice=2838, strikePrice=2820, interestRate=4, dayToExp=27, volatility=None, vol_type='hv2W', cp='C')
    print(p)
